import seaborn as sns
import pandas as pd
import matplotlib.pyplot as plt
from pandas.api.types import CategoricalDtype
from statsmodels.miscmodels.ordinal_model import OrderedModel
import os
from contextlib import redirect_stdout

dir = 'py-replication'
try:
 os.mkdir(dir)
except FileExistsError:
 pass

with open('py-replication/summary.txt', 'w') as f:
	with redirect_stdout(f):
		print("")
		print("************************************")
		print("** File-Name:	replication.py    **")
		print("** Author:		Eddie Hearn       **")
		print("** Purpose:		Hearn 2023 (MSLAS)**")
		print("** Program:		python3           **")
		print("** OS:			Debian GNU/Linux  **")
		print("************************************")
		print("")
		print("The 2020 Presidential Election and Ally Perceptions of the US: A Natural Experiment in Japan")
		print("")
		
df=pd.read_csv('data/py-data.csv')

var = ["us_depend", "us_trust", "us_infl"]
i = 0

while i < 3:
    g =  sns.catplot(y=var[i], data=df, kind='count', hue='elec_tr', palette="pastel",  edgecolor=".6")
    g.set_axis_labels("Number of Respondents", "")
    g.set_yticklabels(["strongly\ndisagree" , "disagree" , "neither" , "agree" , "strongly\nagree"])
    g.set_yticklabels(rotation=20) 
    g._legend.set_title('Experimental Group')
    new_labels = ['Control', 'Treatment']
    for t, l in zip(g._legend.texts, new_labels):
        t.set_text(l)
    sns.move_legend(g, "upper right")
    plt.tight_layout()
    plt.savefig(f'py-replication/{var[i]}.svg')
    i = i + 1

i = 0
while i < 3:
    g =  sns.catplot(x=var[i], data=df, kind='count', hue='elec_tr', col="ldp_s", palette="pastel",  edgecolor=".6")
    titles = ['LDP-K Opposition','Nuetral','LDP-K Suppport']
    for ax,title in zip(g.axes.flatten(),titles):
        ax.set_title(title )
    g._legend.set_title('Experimental Group')
    new_labels = ['Control', 'Treatment']
    for t, l in zip(g._legend.texts, new_labels):
        t.set_text(l)
    
    g.set_axis_labels("", "Number of Respondents")
    g.set_xticklabels(["strongly\ndisagree" , "disagree" , "neither" , "agree" , "strongly\nagree"])
    g.set_xticklabels(rotation=20) 
    plt.savefig(f'py-replication/ldp_{var[i]}.svg')
    i = i + 1




df2 = pd.get_dummies(df['region'], drop_first=True)
df2 = pd.concat([df, df2], axis=1) 
df2 = df2.dropna()

i = 0 
while i < 3:
    cat_type = CategoricalDtype(categories=[1, 2, 3, 4, 5], ordered=True)

    df[var[i]] = df[var[i]].astype(cat_type)
    df[var[i]].dtype
    df2[var[i]] = df2[var[i]].astype(cat_type)
    df2[var[i]].dtype

    basic = OrderedModel(df[var[i]], df[['elec_tr']], distr='probit')
    basic = basic.fit(method='bfgs')
    
    inter = OrderedModel(df[var[i]], df[['elec_tr', 'ldp_s', 'ldp_x_elec']], distr='probit')
    inter = inter.fit(method='bfgs')
    
    controls = OrderedModel(df2[var[i]], df2[['elec_tr', 'ldp_s', 'ldp_x_elec', 'age', 'collgrad', 'inc', 'female', 2, 3, 4, 5, 6, 7, 8]], distr='probit')
    controls = controls.fit(method='bfgs')
    
    with open('py-replication/summary.txt', 'a') as f:
        with redirect_stdout(f):
            print(" ")
            print("*** Basic Model",var[i])
            print(basic.summary().as_text())
            print(" ")
            print(" ")
            print("*** Interaction Model",var[i])
            print(inter.summary().as_text())
            print(" ")
            print(" ")
            print("*** controls",var[i])
            print(controls.summary().as_text())
    i = i + 1

exit(0)
